// Usage: ./ssh-server -repo=/xxx/xxx/xxx/ -port=8883 -hostkey=/etc/ssh/ssh_host_rsa_key
package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	cutils "go-git-protocols/lib"

	"github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
)

var repoRoot *string

func main() {
	repoRoot = flag.String("repo", "/Users/zoker/Tmp/repositories/", "Specify a repositories dir.")
	port := flag.String("port", "8883", "Specify a port to start process.")
	hostKeyPath := flag.String("hostkey", "/etc/ssh/ssh_host_rsa_key", "Specify a private key.")
	flag.Parse()

	// init host key and public key authentication
	var hostOption ssh.Option
	if *hostKeyPath == "" {
		fmt.Println("You must specify a private host key.")
		os.Exit(1)
	} else {
		hostOption = ssh.HostKeyFile(*hostKeyPath)
	}

	keyHandler := func(ctx ssh.Context, key ssh.PublicKey) bool {
		// replace your public key auth logic here
		pubKeyStr := gossh.MarshalAuthorizedKey(key)
		log.Printf("SSH user is %s, using key: %s", ctx.User(), pubKeyStr)
		return true // set false to use password authentication
	}
	keyOption := ssh.PublicKeyAuth(keyHandler)

	// password validate authentication
	pwdHandler := func(ctx ssh.Context, password string) bool {
		// replace your own password auth logic here
		if ctx.User() == "zoker" && password == "zoker" {
			return true
		}
		return false
	}
	pwdOption := ssh.PasswordAuth(pwdHandler)

	// process ssh session pack
	ssh.Handle(func(s ssh.Session) {
		handlePack(s)
	})

	addr := fmt.Sprintf(":%s", *port)
	log.Printf("Starting ssh server on port %s \n", *port)
	log.Fatal(ssh.ListenAndServe(addr, nil, hostOption, pwdOption, keyOption))
}

func handlePack(s ssh.Session) {
	args := s.Command()
	service := args[0]
	repoName := args[1]
	if service != "git-upload-pack" && service != "git-receive-pack" {
		exitSession(s, errors.New("Not allowed command. \n"))
	}

	repoPath := fmt.Sprintf("%s%s", *repoRoot, repoName)
	if err := cutils.CheckRepo(repoPath, service); err != nil {
		exitSession(s, err)
	}

	cmdPack := exec.Command("git", service[4:], repoPath)
	cmdStdin, err := cmdPack.StdinPipe()
	cmdStdout, err := cmdPack.StdoutPipe()
	err = cmdPack.Start()
	if err != nil {
		exitSession(s, err)
	}
	go func() { _, _ = io.Copy(cmdStdin, s) }()
	_, _ = io.Copy(s, cmdStdout)
	_ = cmdPack.Wait()
}

func exitSession(s ssh.Session, err error) {
	errStr := fmt.Sprintf("ERR %s", err.Error())
	pktData := fmt.Sprintf("%04x%s", len(errStr)+4, errStr)
	_, _ = s.Write([]byte(pktData))
	_ = s.Exit(1)
}